#pragma once
#include <iostream>
#include <torch/torch.h>
#include <torch/csrc/api/include/torch/version.h>
#include "NerfStudio.h"
#include "KDTreeTensor.h"
#include "SphericalHarmonics.h"
#include "SSIM.h"
#include "InputData.h"
#include "OptimScheduler.h"

using namespace torch::indexing;
using namespace torch::autograd;

torch::Tensor randomQuatTensor(long long vN);
torch::Tensor projectionMatrix(float vZNear, float vZFar, float vFovX, float vFovY, const torch::Device& vDevice);
torch::Tensor psnr(const torch::Tensor& vRendered, const torch::Tensor& vGroundTruth);
torch::Tensor l1(const torch::Tensor& vRendered, const torch::Tensor& vGroundTruth);

class CModel {
public:
	CModel(const SInputData& vInputData, int vNumCameras,
		int vNumDownscales, int vResolutionSchedule, int vSHDegree, int vSHDegreeInterval,
		int vRefineEvery, int vWarmupLength, int vResetAlphaEvery, float vDensifyGradThresh, float vDensifySizeThresh, float vCullAlphaThresh, int vStopScreenSizeAt, float vSplitScreenSize,
		int vMaxSteps, bool vKeepCrs,
		const torch::Device& vDevice);

	~CModel() {
		delete m_pMeansOpt;
		delete m_pScalesOpt;
		delete m_pQuatsOpt;
		delete m_pFeaturesDcOpt;
		delete m_pFeaturesRestOpt;
		delete m_pOpacitiesOpt;

		delete m_pMeansOptScheduler;
	}

	torch::Tensor forward(CCamera& vCam, int vStep);
	void calPSNRAndL1(std::vector<CCamera>& vCamVec, int vStep, float vSSIMWeight, torch::Device vDevice);
	void optimizersZeroGrad();
	void optimizersStep();
	void schedulersStep(int vStep);
	int getDownscaleFactor(int vStep);
	void afterTrain(int vStep);
	void saveScene(std::filesystem::path& vOutputPath, std::string& vSceneName, int vStep, SInputData* vInputData, std::vector<CCamera>& vCams);
	void savePly(const std::string& vFileName);
	void saveSplat(const std::string& vFileName);
	void saveDebugPly(const std::string& vFileName);
	torch::Tensor mainLoss(torch::Tensor& vRGB, torch::Tensor& vGroudTruth, float vSSIMWeight);

	void addToOptimizer(torch::optim::Adam* vOptimizer, const torch::Tensor& vNewParam, const torch::Tensor& vIdcs, int vNSamples);
	void removeFromOptimizer(torch::optim::Adam* vOptimizer, const torch::Tensor& vNewParam, const torch::Tensor& vDeletedMask);

private:
	torch::Tensor m_Means;
	torch::Tensor m_Scales;
	torch::Tensor m_Quats;
	torch::Tensor m_FeaturesDc;
	torch::Tensor m_FeaturesRest;
	torch::Tensor m_Opacities;

	torch::optim::Adam* m_pMeansOpt;
	torch::optim::Adam* m_pScalesOpt;
	torch::optim::Adam* m_pQuatsOpt;
	torch::optim::Adam* m_pFeaturesDcOpt;
	torch::optim::Adam* m_pFeaturesRestOpt;
	torch::optim::Adam* m_pOpacitiesOpt;

	COptimScheduler* m_pMeansOptScheduler;

	torch::Tensor m_Radii; // set in forward()
	torch::Tensor m_XYs; // set in forward()
	int m_LastHeight; // set in forward()
	int m_LastWidth; // set in forward()

	torch::Tensor m_XYsgradNorm; // set in afterTrain()
	torch::Tensor m_VisCounts; // set in afterTrain()  
	torch::Tensor m_Max2DSize; // set in afterTrain()

	float m_SSIMWeight;
	torch::Tensor m_BackgroundColor;
	torch::Device m_Device;
	CSSIM m_SSIM;

	int m_NumCameras;
	int m_NumDownscales;
	int m_ResolutionSchedule;
	int m_SHDegree;
	int m_SHDegreeInterval;
	int m_RefineEvery;
	int m_WarmupLength;
	int m_ResetAlphaEvery;
	float m_CullAlphaThresh;
	int m_StopSplitAt;
	float m_DensifyGradThresh;
	float m_DensifySizeThresh;
	int m_StopScreenSizeAt;
	float m_SplitScreenSize;
	int m_MaxSteps;
	bool m_KeepCrs;

	float m_Scale;
	torch::Tensor m_Translation;
};